#!./bin/python3.6
# -*- coding:utf-8 -*-

from rknn.api import RKNN

INPUT_SIZE = 64

if __name__ == '__main__':
    # 创建RKNN执行对象
    rknn = RKNN()
    # 配置模型输入，用于NPU对数据输入的预处理
    # channel_mean_value='0 0 0 255'，那么模型推理时，将会对RGB数据做如下转换
    # (R - 0)/255, (G - 0)/255, (B - 0)/255。推理时，RKNN模型会自动做均值和归一化处理
    # reorder_channel=’0 1 2’用于指定是否调整图像通道顺序，设置成0 1 2即按输入的图像通道顺序不做调整
    # reorder_channel=’2 1 0’表示交换0和2通道，如果输入是RGB，将会被调整为BGR。如果是BGR将会被调整为RGB
    # 图像通道顺序不做调整
    # rknn.config(channel_mean_value='0 0 0 255', reorder_channel='0 1 2')
    # rknn.config(
    #     mean_values=[[123.675, 116.28, 103.53]],
    #     std_values=[[58.395, 58.395, 58.395]],
    #     reorder_channel='0 1 2')
    add_perm = False  # 如果设置成True,则将模型输入layout修改成NHWC
    rknn.config(
        batch_size=1,
        mean_values=[[0, 0, 0]],
        std_values=[[255, 255, 255]],
        reorder_channel='0 1 2',
        force_builtin_perm=add_perm,
        output_optimize=1,
        target_platform='rv1126')

    # 加载TensorFlow模型
    # tf_pb='digital_gesture.pb'指定待转换的TensorFlow模型
    # inputs指定模型中的输入节点
    # outputs指定模型中输出节点
    # input_size_list指定模型输入的大小
    print('--> Loading model')
    # rknn.load_tensorflow(tf_pb='digital_gesture.pb',
    #                      inputs=['input_x'],
    #                      outputs=['probability'],
    #                      input_size_list=[[INPUT_SIZE, INPUT_SIZE, 3]])

    # rknn.load_onnx("/home/roota/Desktop/AI/rknntools/yolov5s.onnx")
    # print('done')
    ret = rknn.load_onnx(model='best.onnx',
                         inputs=['images'],
                         outputs=['output80x', 'output40x', 'output20x'],
                         input_size_list=[[3, 640, 640]], )
    if ret != 0:
        print('Load  onnx  failed!')
        exit(ret)
    print('done')

    # 创建解析pb模型
    # do_quantization=False指定不进行量化
    # 量化会减小模型的体积和提升运算速度，但是会有精度的丢失
    print('--> Building model')
    rknn.build(do_quantization=True)
    print('done')

    # 导出保存rknn模型文件
    ret = rknn.export_rknn('./best.rknn')
    if ret != 0:
        print('Export  rknn failed!')
        exit(ret)
    else:
        print('Export  rknn success!')

    # Release RKNN Context
    rknn.release()
